import pathlib
import pickle
import unittest
from datetime import datetime, date
import pandas as pd
import numpy as np
from tqdm import tqdm
from matplotlib.dates import date2num
# Need to install the quant package itself first (pip install -e .)
from quant import strategy, utils
from quant.strategy import TradeOp
from quant.backtest import backtest
from ruamel import yaml
from easydict import EasyDict
from typing import List


class TestGrid(unittest.TestCase):
    testdata = pathlib.Path(__file__).parent.joinpath('testdata')
    with pathlib.Path(__file__).parent.joinpath('backtest/backtest.yml').open(
            mode='r') as f:
        docs = {doc['name']: EasyDict(doc) for doc in yaml.safe_load_all(f)}

    def grid_test(self, signals: pd.DataFrame, name: str, grid: List[float],
            frequency: str):
        f = self.testdata.joinpath('%s.csv' % name)
        if not f.exists():
            signals['signal'] = signals.signal.astype(int)
            signals.to_csv(f)
            f = self.testdata.joinpath('%s.png' % name)
            fig = utils.plot_trade_signals(signals.reset_index(), 'datetime', 'open',
                                           frequency, grid)
            fig.savefig(self.testdata.joinpath('%s.png' % name))
        else:
            expected = pd.read_csv(f,
                                   parse_dates=['datetime'],
                                   index_col=['datetime'])
            np.testing.assert_array_equal(
                expected.signal.tolist(), signals.signal.tolist())
            np.testing.assert_array_equal(
                expected.grid_id.tolist(), signals.grid_id.tolist())

    def test_basic_grid(self):
        name = 'BasicGrid'
        bt = EasyDict(self.docs['grid_dfcf_1d'].copy())
        bt.start_date = '2021-01-01'
        bt.end_date = '2022-01-01'
        bt.plot = False
        bt.grid = 'strategy.grid_by_lh(26, 40, 10)'
        bt.strategies = [bt.strategies[0]]
        results = backtest.runtest(bt)
        signals = results[name]
        if bt.cache_dir is None:
            signals['datetime'] = pd.to_datetime(signals.datetime)
            signals.set_index('datetime', inplace=True)
        self.grid_test(signals, name, eval(bt.grid).tolist(), bt.frequency)

    def test_enhanced_basic_grid(self):
        name = 'EnhancedBasicGrid'
        bt = EasyDict(self.docs['grid_dfcf_1d'].copy())
        bt.start_date = '2021-01-01'
        bt.end_date = '2022-01-01'
        bt.plot = False
        bt.grid = 'strategy.grid_by_lh(26, 40, 10)'
        bt.strategies = [bt.strategies[1]]
        results = backtest.runtest(bt)
        signals = results[name]
        if bt.cache_dir is None:
            signals['datetime'] = pd.to_datetime(signals.datetime)
            signals.set_index('datetime', inplace=True)
        self.grid_test(signals, name, eval(bt.grid).tolist(), bt.frequency)

    def test_basic_dynamic_grid(self):
        name = 'BasicDynamicGrid'
        bt = EasyDict(self.docs['grid_dfcf_1d'].copy())
        bt.start_date = '2021-01-01'
        bt.end_date = '2022-01-01'
        bt.plot = False
        bt.grid = 'strategy.grid_by_lh(26, 40, 10)'
        bt.strategies = [bt.strategies[2]]
        results = backtest.runtest(bt)
        signals = results[name]
        if bt.cache_dir is None:
            signals['datetime'] = pd.to_datetime(signals.datetime)
            signals.set_index('datetime', inplace=True)
        self.grid_test(signals, name, eval(bt.grid).tolist(), bt.frequency)

    def test_static_turning_point_grid(self):
        name = 'StaticTurningPointGrid'
        bt = EasyDict(self.docs['grid_dfcf_1m'].copy())
        bt.start_date = '2022-01-01'
        bt.end_date = '2022-01-10'
        bt.plot = False
        bt.grid = 'strategy.grid_by_mid_step_percent(35, 1.0 / 100, 16)'
        bt.strategies = [bt.strategies[0]]
        bt.strategies[0]['strategy'] = name
        results = backtest.runtest(bt)
        signals = results[name]
        if bt.cache_dir is None:
            signals['datetime'] = pd.to_datetime(signals.datetime)
            signals.set_index('datetime', inplace=True)
        self.grid_test(signals, name, eval(bt.grid).tolist(), bt.frequency)

    def test_dynamic_turning_point_grid(self):
        name = 'DynamicTurningPointGrid'
        bt = EasyDict(self.docs['grid_dfcf_1m'].copy())
        bt.start_date = '2022-01-01'
        bt.end_date = '2022-01-10'
        bt.plot = False
        bt.grid = 'strategy.grid_by_mid_step_percent(35, 1.0 / 100, 16)'
        bt.strategies = [bt.strategies[1]]
        results = backtest.runtest(bt)
        signals = results[name]
        if bt.cache_dir is None:
            signals['datetime'] = pd.to_datetime(signals.datetime)
            signals.set_index('datetime', inplace=True)
        self.grid_test(signals, name, eval(bt.grid).tolist(), bt.frequency)

    def test_grid_by_fn(self):
        actual = strategy.grid_by_mid_step(5, 1, 10)
        np.testing.assert_array_equal(
            [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0], actual)

        actual = strategy.grid_by_mid_step_percent(5, 0.2, 10)
        np.testing.assert_array_almost_equal(
            [2.048, 2.56, 3.2, 4.0, 5.0, 6.0, 7.2, 8.64, 10.368, 12.4416],
            actual)

        actual = strategy.grid_by_lh(1, 10, 10)
        np.testing.assert_array_equal(
            [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0], actual)


class TestConditionTrade(unittest.TestCase):
    def test_turning_point_buy(self):
        st = strategy.TurningPointBuy(threshold=1)
        testdata = pd.DataFrame.from_dict(
            {'open': [10, 9.6, 9.3, 9.5, 9.7, 9.2, 9.8, 10.3, 10.4, 9.9]})
        testdata['bob'] = pd.date_range(start=datetime.now(),
                                        periods=len(testdata),
                                        freq='m')
        expected = [
            TradeOp.SKIP_BUY, TradeOp.SKIP_BUY, TradeOp.SKIP_BUY,
            TradeOp.SKIP_BUY, TradeOp.SKIP_BUY, TradeOp.SKIP_BUY,
            TradeOp.SKIP_BUY, TradeOp.BUY, TradeOp.BUY, TradeOp.SKIP_BUY
        ]
        for idx in testdata.index:
            dk = testdata.loc[idx]
            actual = st.generate_signal(pd.DataFrame([dk]),
                                        datetime.now()).iloc[0]
            self.assertEqual(actual.signal, expected[idx])

    def test_turning_point_sell(self):
        st = strategy.TurningPointSell(threshold=1)
        testdata = pd.DataFrame.from_dict(
            {'open': [10, 9.6, 10.3, 9.9, 10.6, 10, 9.8, 9.5, 9.3, 10]})
        testdata['bob'] = pd.date_range(start=datetime.now(),
                                        periods=len(testdata),
                                        freq='m')
        expected = [
            TradeOp.SKIP_SELL, TradeOp.SKIP_SELL, TradeOp.SKIP_SELL,
            TradeOp.SKIP_SELL, TradeOp.SKIP_SELL, TradeOp.SKIP_SELL,
            TradeOp.SKIP_SELL, TradeOp.SELL, TradeOp.SELL, TradeOp.SKIP_SELL
        ]
        for idx in testdata.index:
            dk = testdata.loc[idx]
            actual = st.generate_signal(pd.DataFrame([dk]),
                                        datetime.now()).iloc[0]
            self.assertEqual(actual.signal, expected[idx])


class TestConbond(unittest.TestCase):
    testdata = pathlib.Path(__file__).parent.joinpath('testdata')
    with pathlib.Path(__file__).parent.joinpath('backtest/backtest.yml').open(
            mode='r') as f:
        docs = {doc['name']: EasyDict(doc) for doc in yaml.safe_load_all(f)}

    def conbond_test(self, signals: pd.DataFrame, name: str):
        f = self.testdata.joinpath('%s.csv' % name)
        if not f.exists():
            signals.to_csv(f)
        else:
            expected = pd.read_csv(f,
                                   parse_dates=['datetime'],
                                   index_col=['datetime', 'order_book_id'])
            pd.testing.assert_frame_equal(
                expected.reset_index()[['order_book_id']],
                signals.reset_index()[['order_book_id']])

    def test_double_low(self):
        name = 'double_low'
        bt = self.docs['conbond']
        bt.start_date = '2021-06-01'
        bt.end_date = '2021-07-01'
        bt.plot = False
        bt.strategies = [{
            'strategy': 'ConbondRotateStrategy',
            'params': {
                'name': name,
                'bars': None,
                'instruments': None
            }
        }]
        results = backtest.runtest(bt)
        signals = results[name]
        self.conbond_test(signals, name)


if __name__ == '__main__':
    unittest.main()
